package covertjava.classloader;

import java.util.*;
import java.io.*;
import java.util.zip.*;

/**
* <p>Custom class loader that delegates to a decorator for runtime bytecode instrumentation</p>
* <p>Copyright: Copyright (c) 2004 Sams Publishing</p>
* @author Alex Kalinovsky
* @version 1.0
*/
public class DecoratingClassLoader extends ClassLoader {

    protected ClassDecorator decorator;
    protected ArrayList classPathItems = new ArrayList();

    public DecoratingClassLoader() {
        this(System.getProperty("decorate.class.path"));
    }

    public DecoratingClassLoader(String classPath) {
        parseClassPath(classPath);
    }

    public ClassDecorator getDecorator() {
        return this.decorator;
    }

    public void setDecorator(ClassDecorator decorator) {
        this.decorator = decorator;
    }

    protected void parseClassPath(String classPath) {
        if (classPath == null)
            throw new IllegalArgumentException("DecoratingClassLoader must be initialized with class path (\"decorate.class.path\" system property is used by default)");

        String separator = System.getProperty("path.separator");
        StringTokenizer pathTokenizer = new StringTokenizer(classPath, separator);
        while (pathTokenizer.hasMoreTokens()) {
            classPathItems.add(pathTokenizer.nextToken());
        }
    }

    protected Class findClass(String name) throws ClassNotFoundException {

        System.out.println("DecoratingClassLoader loading class " + name);

        byte[] bytes = getClassBytes(name);

        if (getDecorator() != null)
            bytes = getDecorator().decorateClass(name, bytes);

        return defineClass(name, bytes, 0, bytes.length);
    }

    protected byte[] getClassBytes(String className) throws ClassNotFoundException {
        byte[] bytes = null;
        for (int i= 0; i < classPathItems.size(); i++) {
            String path = (String)classPathItems.get(i);
            String fileName = className.replace('.', '/') + ".class";
            if (isJar(path) == true) {
                bytes = loadJarBytes(path, fileName);
            }
            else {
                bytes = loadFileBytes(path, fileName);
            }
            if (bytes != null)
                break;
        }

        if (bytes == null)
            throw new ClassNotFoundException(className);

        return bytes;
    }

    boolean isJar(String pathEntry) {
        return pathEntry.endsWith(".jar") || pathEntry.endsWith(".zip");
    }

    protected byte[] loadFileBytes(String path, String fileName) {
        File file= new File(path, fileName);
        if (file.exists()) {
            return loadClassBytes(file);
        }
        return null;
    }

    protected byte[] loadClassBytes(File f) {
        try {
            FileInputStream stream = new FileInputStream(f);
            ByteArrayOutputStream out = new ByteArrayOutputStream(1000);
            byte[] b = new byte[1000];
            int n;
            while ((n= stream.read(b)) != -1)
                out.write(b, 0, n);
            stream.close();
            out.close();
            return out.toByteArray();

        } catch (IOException e) {
        }
        return null;
    }

    protected byte[] loadJarBytes(String path, String fileName) {
        ZipFile zipFile = null;
        InputStream stream = null;
        File archive = new File(path);
        if (!archive.exists())
            return null;

        try {
            zipFile= new ZipFile(archive);
        } catch(IOException io) {
            return null;
        }

        ZipEntry entry= zipFile.getEntry(fileName);
        if (entry == null)
            return null;

        int size = (int) entry.getSize();
        try {
            stream = zipFile.getInputStream(entry);
            byte[] bytes = new byte[size];
            int pos = 0;
            while (pos < size) {
                int n= stream.read(bytes, pos, bytes.length - pos);
                pos += n;
            }
            zipFile.close();
            return bytes;
        }
        catch (IOException e) {
        }
        finally {
            try {
                if (stream != null)
                    stream.close();
            } catch (IOException e) {
            }
        }
        return null;
    }

}
